#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# Copyright (c) 2020, 2024, Niklas Hauser
#
# This file is part of the modm project.
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
# -----------------------------------------------------------------------------

from pathlib import Path
from collections import defaultdict
tusb_config = {}
default_port = 0

# -----------------------------------------------------------------------------
def init(module):
    module.name = ":tinyusb"
    module.description = FileReader("module.md")

def prepare(module, options):
    device = options[":target"]
    has_otg_hs = device.has_driver("usb_otg_hs")
    has_otg_fs = device.has_driver("usb_otg_fs")
    if not (device.has_driver("usb") or
            has_otg_fs or
            has_otg_hs or
            device.has_driver("udp") or
            device.has_driver("uhp")):
        return False

    configs = ["device.cdc", "device.dfu_rt", "device.midi", "device.msc", "device.vendor"]
    # TODO: Allow all device classes
    # configs  = {"device.{}".format(p.parent.name)
    #            for p in Path(localpath("tinyusb/src/class/")).glob("*/*_device.h")}
    # TODO: Allow all host classes
    # configs |= {"host.{}".format(p.parent.name)
    #            for p in Path(localpath("tinyusb/src/class/")).glob("*/*_host.h")}

    module.add_list_option(
            EnumerationOption(name="config",
                              description="Endpoint Configuration",
                              enumeration=configs,
                              dependencies=lambda vs:
                                    [":tinyusb:" + v.replace(".", ":") for v in vs]))


    # Most devices only have FS=usb, some STM32 have FS/HS=usb_otg_fs/usb_otg_hs
    # and some STM32 only have HS=usb_otg_hs, so we only define this option if
    # there is a choice and default this to FS by default.
    if has_otg_hs:
        module.add_option(
                EnumerationOption(name="max-speed", description="Maximum HS port speed",
                                  enumeration=["full", "high"], default="high"))
        # DEPRECATED: 2025q1
        module.add_alias(
                Alias(name="speed", description="Renamed for more clarity",
                      destination=":tinyusb:max-speed"))
    # Detect the number of ports and the default port
    ports = 2 if (has_otg_hs and has_otg_fs) else 1
    if has_otg_hs and not has_otg_fs:
        global default_port
        default_port = 1
    # Generate the host and device modules
    module.add_submodule(TinyUsbDeviceModule(ports))
    module.add_submodule(TinyUsbHostModule(ports))
    module.depends(":cmsis:device", ":architecture:atomic", ":architecture:interrupt", ":platform:usb")
    return True


def validate(env):
    has_device = env.has_module(":tinyusb:device")
    has_host = env.has_module(":tinyusb:host")
    if has_device and has_host:
        device_port = env.get(":tinyusb:device:port")
        host_port = env.get(":tinyusb:host:port")
        if device_port == host_port and host_port is not None:
            raise ValidateException("You cannot place Device and Host on the same USB port!")


def build(env):
    env.collect(":build:path.include", "modm/ext/tinyusb/")
    env.outbasepath = "modm/ext/tinyusb"

    env.copy("tinyusb/src/tusb.h", "tusb.h")
    env.copy("tinyusb/src/tusb_option.h", "tusb_option.h")
    env.copy("tinyusb/src/osal/osal.h", "osal/osal.h")
    env.copy("tinyusb/src/common", "common/")
    env.copy("tinyusb/src/tusb.c", "tusb.c")

    """ Generic config defaults:
    - CFG_TUSB_CONFIG_FILE  = tusb_config.h
    - CFG_TUSB_DEBUG        = 0
    - CFG_TUSB_DEBUG_PRINTF = [undef]
    - CFG_TUSB_MCU          = OPT_MCU_NONE
    - CFG_TUSB_MEM_ALIGN    = __attribute__((aligned(4)))
    - CFG_TUSB_MEM_SECTION  = [empty] # FIXME
    - CFG_TUSB_OS           = OPT_OS_NONE
    - CFG_TUSB_RHPORT0_MODE = OPT_MODE_NONE
    - CFG_TUSB_RHPORT1_MODE = OPT_MODE_NONE
    """
    global tusb_config

    target = env[":target"].identifier
    if target.platform == "stm32":
        tusb_config["CFG_TUSB_MCU"] = f"OPT_MCU_STM32{target.family.upper()}"
        # TODO: use modm-devices driver type for this
        fs_dev = (target.family in ["f0", "f3", "l0", "g4"] or
                 (target.family == "f1" and target.name <= "03"))
        if fs_dev:
            # PMA buffer size: 512B or 1024B
            env.copy("tinyusb/src/portable/st/stm32_fsdev/", "portable/st/stm32_fsdev/")
        else:
            env.copy("tinyusb/src/portable/synopsys/dwc2/", "portable/synopsys/dwc2/")

    elif target.platform == "sam":
        if target.family == "g5x":
            tusb_config["CFG_TUSB_MCU"] = "OPT_MCU_SAMG"
            env.copy("tinyusb/src/portable/microchip/samg/", "portable/microchip/samg/")
        else:
            series = "e5x" if target.series.startswith("e5") else target.series
            tusb_config["CFG_TUSB_MCU"] = f"OPT_MCU_SAM{series.upper()}"
            env.copy("tinyusb/src/portable/microchip/samd/", "portable/microchip/samd/")

    elif target.platform == "rp":
        if target.family == "20":
            tusb_config["CFG_TUSB_MCU"] = "OPT_MCU_RP2040"
            tusb_config["TUD_OPT_RP2040_USB_DEVICE_UFRAME_FIX"] = "0"
            env.copy("tinyusb/src/portable/raspberrypi/rp2040/", "portable/raspberrypi/rp2040/")

    # TinyUSB has a OS abstraction layer for FreeRTOS, but it is causes problems with modm
    # if env.has_module(":freertos"):
    #    tusb_config["CFG_TUSB_OS"] = "OPT_OS_FREERTOS"
    #    env.copy("tinyusb/src/osal/osal_freertos.h", "osal/osal_freertos.h")

    if env.has_module(":freertos"):
        env.log.warning("TinyUSB in modm does not currently use the FreeRTOS abstraction layer"
                        " and is not thread-safe with FreeRTOS threads.")

    env.copy("tinyusb/src/osal/osal_none.h", "osal/osal_none.h")

    if env.has_module(":debug"):
        tusb_config["CFG_TUSB_DEBUG_PRINTF"] = "tinyusb_debug_printf"
        # env.collect(":build:cppdefines.debug", "CFG_TUSB_DEBUG=2")

    ports = {}
    global default_port
    for mode in ["device", "host"]:
        if env.has_module(f":tinyusb:{mode}"):
            port = env.get(f":tinyusb:{mode}:port", default_port)
            ports[port] = mode[0]
            mode = f"OPT_MODE_{mode.upper()}"
            if port == 1 and (speed := env.get("max-speed")) is not None:
                mode = f"({mode} | OPT_MODE_{speed.upper()}_SPEED)"
            tusb_config[f"CFG_TUSB_RHPORT{port}_MODE"] = mode

    itf_config = env["config"]
    # Enumerate the configurations
    config_enum = defaultdict(list)
    config_enum_counter = 0
    for conf in itf_config:
        config_enum[conf.split(".")[1]].append(config_enum_counter)
        config_enum_counter += 1

    # Generate the ITF and Endpoint counters
    itfs = []
    itf_enum = []
    endpoints = {}
    endpoint_counter = 0
    for devclass, devitfs in config_enum.items():
        prefix = devclass.upper()
        for itf in devitfs:
            endpoint_counter += 1
            itf_prefix = prefix + str(itf)
            itfs.append( (prefix, itf_prefix) )
            if devclass in ["cdc"]:
                itf_enum.extend([itf_prefix, itf_prefix + "_DATA"])
                endpoints[itf_prefix + "_NOTIF"] = hex(0x80 | endpoint_counter)
                endpoint_counter += 1
            elif devclass in ["msc", "midi", "vendor", "dfu_rt"]:
                itf_enum.append(itf_prefix)
            else:
                raise ValidateException(f"Unknown ITF device class '{devclass}'!")

            endpoints[itf_prefix + "_OUT"] = hex(endpoint_counter)
            # SAMG doesn't support using the same EP number for IN and OUT
            if target.platform == "sam" and target.family in ["g"]:
                endpoint_counter += 1
            endpoints[itf_prefix + "_IN"] = hex(0x80 | endpoint_counter)

    if target.platform == "stm32":
        irqs = []
        for port, uirqs in env.query(":platform:usb:irqs")["port_irqs"].items():
            port = {"fs": 0, "hs": 1}[port]
            if port in ports:
                irqs.extend([(irq, port, ports[port]) for irq in uirqs])
    elif target.platform == "sam" and target.family in ["g5x"]:
        irqs = [("UDP", 0, "d")]
    elif target.platform == "sam" and target.family in ["d5x/e5x"]:
        irqs = [("USB_OTHER", 0, "d"), ("USB_SOF_HSOF", 0, "d"),
                ("USB_TRCPT0", 0, "d"), ("USB_TRCPT1", 0, "d")]
    elif target.platform == "rp" and target.family in ["20"]:
        irqs = [("USBCTRL_IRQ", 0, "d")]
    else:
        irqs = [("USB", 0, "d")]

    env.substitutions = {
        "target": target,
        "config": tusb_config,
        "irqs": irqs,
        "ports": ports,
        "with_debug": env.has_module(":debug"),
        "with_cdc": env.has_module(":tinyusb:device:cdc"),
        "itfs": itfs,
        "itfs_enum": itf_enum,
        "endpoints": endpoints,
    }
    env.template("tusb_config.h.in")
    if len(itf_config): env.template("tusb_descriptors.c.in");
    env.template("tusb_port.cpp.in")

    # Add support files shared between device and host classes
    classes = env.generated_local_files(filterfunc=lambda path: "_device.c" in path or "_host.c" in path)
    classes = {(Path(c).parent.name, Path(c).name.rsplit("_", 1)[0]) for c in classes}
    # Add support files outside of naming scheme
    if any(p == "net" for p, _ in classes): classes.update({("net", "ncm"), ("net", "net_device")})
    if ("dfu", "dfu_rt") in classes: classes.add(("dfu", "dfu"))
    # Filter out classes without support files
    classes = ((p, n) for p, n in classes if n not in ["bth", "dfu_rt", "vendor", "ecm_rndis", "ecm"])
    # Now copy the shared support file
    for parent, name in classes:
        env.copy(f"tinyusb/src/class/{parent}/{name}.h", f"class/{parent}/{name}.h")


# -----------------------------------------------------------------------------
class TinyUsbDeviceModule(Module):
    def __init__(self, ports):
        self.ports = ports

    def init(self, module):
        module.name = ":tinyusb:device"
        module.description = """
# TinyUSB in Device Mode

Configuration options:

- `CFG_TUD_ENDPOINT0_SIZE` = 64
- `CFG_TUD_EP_MAX` = 9
- `CFG_TUD_TASK_QUEUE_SZ` = 16
"""

    def prepare(self, module, options):
        classes = {(p.parent.name, p.name.replace("_device.c", ""))
                   for p in Path(localpath("tinyusb/src/class/")).glob("*/*_device.c")}
        for parent, name in classes:
            module.add_submodule(TinyUsbClassModule(parent, name, "device"))
        if self.ports == 2:
            module.add_option(
                    EnumerationOption(name="port", description="USB Port selection",
                                      enumeration={"fs": 0, "hs": 1}, default="fs",
                                      dependencies=lambda s: f":platform:usb:{s[0]}s"))
        return True

    # def validate(self, env):
    #     if env.has_module(":tinyusb:device:ncm") and env.has_module(":tinyusb:device:ecm_rndis"):
    #         raise ValidateException("TinyUSB cannot enable both ECM_RNDIS and NCM network drivers!")

    def build(self, env):
        env.outbasepath = "modm/ext/tinyusb"
        env.copy("tinyusb/src/device/", "device/")


# -----------------------------------------------------------------------------
class TinyUsbHostModule(Module):
    def __init__(self, ports):
        self.config = {}
        self.ports = ports

    def init(self, module):
        module.name = ":tinyusb:host"
        module.description = """
# TinyUSB in Host Mode

Configuration options:

- `CFG_TUH_EP_MAX` = 9
- `CFG_TUH_TASK_QUEUE_SZ` = 16
- `CFG_TUH_ENUMERATION_BUFSZIE` [sic] = 256
"""

    def prepare(self, module, options):
        # On STM32, only OTG has Host Mode
        if not (options[":target"].has_driver("usb_otg_fs") or
                options[":target"].has_driver("usb_otg_hs") or
                options[":target"].has_driver("uhp:samg*")):
            return False

        classes = {(p.parent.name, p.name.replace("_host.c", ""))
                   for p in Path(localpath("tinyusb/src/class/")).glob("*/*_host.c")}
        for parent, name in classes:
            module.add_submodule(TinyUsbClassModule(parent, name, "host"))
        if self.ports == 2:
            module.add_option(
                    EnumerationOption(name="port", description="USB Port selection",
                                      enumeration={"fs": 0, "hs": 1}, default="hs",
                                      dependencies=lambda s: f":platform:usb:{s[0]}s"))
        return True

    def build(self, env):
        env.outbasepath = "modm/ext/tinyusb"
        env.copy("tinyusb/src/host/", "host/")


# -----------------------------------------------------------------------------
# Only contains those configurations that this modules can generate a descriptor for
tu_config_descr = {"device": {
    "bth": """
- `CFG_TUD_BTH_ISO_ALT_COUNT` = [undef] modm default: 1
""",
    "cdc": """
- `CFG_TUD_CDC_EP_BUFSIZE` = 64/512 (fs/hs)
- `CFG_TUD_CDC_RX_BUFSIZE` = [undef] modm default: 512
- `CFG_TUD_CDC_TX_BUFSIZE` = [undef] modm default: 512
""",
    "midi": """
- `CFG_TUD_MIDI_EP_BUFSIZE` = 64/512 (fs/hs)
- `CFG_TUD_MIDI_RX_BUFSIZE` = [undef] modm default: 64/512 (fs/hs)
- `CFG_TUD_MIDI_TX_BUFSIZE` = [undef] modm default: 64/512 (fs/hs)
""",
    "msc": """
- `CFG_TUD_MSC_EP_BUFSIZE` = [undef] modm default: 512
""",
    "vendor": """
- `CFG_TUD_VENDOR_EPSIZE` = 64
- `CFG_TUD_VENDOR_RX_BUFSIZE` = [undef] modm default: 64/512 (fs/hs)
- `CFG_TUD_VENDOR_TX_BUFSIZE` = [undef] modm default: 64/512 (fs/hs)
"""
}
}
class TinyUsbClassModule(Module):
    def __init__(self, parent, name, mode):
        self.parent = parent
        self.name = name
        self.mode = mode

    def init(self, module):
        module.name = f":tinyusb:{self.mode}:{self.name}"
        descr = f"# {self.mode.capitalize()} class {self.name.upper()}"
        conf = tu_config_descr.get(self.mode, {}).get(self.name, "")
        if conf: descr += "\n\nConfiguration options:\n" + conf
        else: descr += "\n\nPlease consult the TinyUSB docs for configuration options.\n"
        module.description = descr

    def prepare(self, module, options):
        if self.mode == "device":
            if self.name == "cdc":
                module.depends(":architecture:uart")
            if self.name == "midi":
                module.depends(":tinyusb:device:audio")
        module.depends(":tinyusb")
        return True

    def build(self, env):
        env.outbasepath = "modm/ext/tinyusb"
        for suffix in ["h", "c"]:
            # both classes share the `net_device.h` file
            if self.name in ["ncm", "ecm_rndis"] and suffix == "h":
                continue
            file = f"class/{self.parent}/{self.name}_{self.mode}.{suffix}"
            env.copy(f"tinyusb/src/{file}", file)

        global tusb_config
        cfg_name = f"CFG_TU{self.mode[0].upper()}_{self.name.upper()}"
        if "DFU_RT" in cfg_name: cfg_name = cfg_name.replace("_RT", "_RUNTIME")
        tusb_config[cfg_name] = env[":tinyusb:config"].count(f"{self.mode}.{self.name}")
        is_hs = env.get(":tinyusb:max-speed", "full") == "high"

        if self.mode == "device":
            # These are the defaults that don't crash TinyUSB.
            if self.name in ["cdc", "midi", "vendor"]:
                tusb_config[cfg_name+"_RX_BUFSIZE"] = 512 if is_hs else 64
                tusb_config[cfg_name+"_TX_BUFSIZE"] = 512 if is_hs else 64
            if self.name in ["msc"]:
                tusb_config[cfg_name+"_EP_BUFSIZE"] = 512
            if self.name in ["bth"]:
                tusb_config[cfg_name+"_ISO_ALT_COUNT"] = 1
